import torch.nn as nn
import torch
from torch.utils import data
from config.Load_DefaultConfig import DefaultConfig
from models.backbone import model
from dataset.dataset import CrackDataSet
from tqdm import tqdm
import torchnet
import os.path as path


if __name__ == '__main__':
    opt = DefaultConfig()
    model = model()
    dataset = CrackDataSet(opt.data_set)
    data_loader = data.DataLoader(
        dataset,
        batch_size=opt.batch_size,
        shuffle=opt.shuffle,
        num_workers=opt.num_workers
    )

    device = torch.device(opt.device if torch.cuda.is_available() else "cpu")   # 模型

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=opt.lr, weight_decay=opt.weight_decay)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, "min", factor=opt.factor, patience=opt.patience, verbose=True)
    model.to(device)
    loss_meter = torchnet.meter.AverageValueMeter()
    for epoch in range(opt.max_epoch):
        loss_meter.reset()
        tqdm_bar = tqdm(enumerate(data_loader), desc=f"epoch{epoch}", total=len(data_loader))
        for ii, (img, label) in tqdm_bar:
            optimizer.zero_grad()
            img = img.to(device)
            label = label.to(device)
            out_label = model(img)
            loss = criterion(out_label, label)
            loss_meter.add(float(loss.data))
            loss.backward()  # 反向传播
            optimizer.step()  # 更新优化器
            tqdm_bar.set_postfix_str({
                "loss_now": f"{float(loss.data):.5f}",
                "loss_mean": f"{float(loss_meter.mean):.5f}"
            })
        scheduler.step(loss_meter.mean)
        torch.save(model, path.join(opt.model_path, f"epoch{epoch}.pth"))
